import os
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import pandas as pd
import yaml


def load_config(config_path):
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)


def getDef():
    project_path = os.environ.get("PROJECT_PATH")
    config_path = os.path.join(project_path, 'config', 'ssc_config.yaml')
    config = load_config(config_path)
    df = pd.read_csv(config["data"]["processed_data_path"])
    return df


# 获取每一期  奇数和偶数的个数
def getOddEven(row):
    numbers = [row['num1'], row['num2'], row['num3'], row['num4'], row['num5']]
    odd_count = sum(1 for n in numbers if n % 2 != 0)  # 奇数个数
    even_count = 5 - odd_count  # 偶数个数
    return odd_count, even_count


# 把每期奇偶数的数量封装进 df
def addOddEven(df):
    df['odd_count'], df['even_count'] = zip(*df.apply(getOddEven, axis=1))
    return df


# 获取最近几期的奇数和偶数的平均值，作为新特征
def getCurrentAvgOddEven(df):
    project_path = os.environ.get("PROJECT_PATH")
    config_path = os.path.join(project_path, 'config', 'ssc_config.yaml')
    config = load_config(config_path)
    df['rolling_odd_mean'] = df['odd_count'].rolling(window=config["data"]["window_size"]).mean()
    df['rolling_even_mean'] = df['even_count'].rolling(window=config["data"]["window_size"]).mean()
    # 填补缺失值
    df.fillna(0, inplace=True)
    return df


# 定于目标变量： 奇数是否大于偶数
def setTarget(df):
    df['target'] = df['odd_count'].apply(lambda x: 1 if x > 2 else 0)
    return df


# 加入训练集和测试集
def addTrainTest(df):
    project_path = os.environ.get("PROJECT_PATH")
    config_path = os.path.join(project_path, 'config', 'ssc_config.yaml')
    config = load_config(config_path)
    # 特征列
    features = ['odd_count', 'even_count', 'rolling_odd_mean', 'rolling_even_mean']

    # 特征数据和目标数据
    X = df[features]
    y = df['target']

    # 划分训练集和测试集
    x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=config["data"]["test_size"],
                                                        random_state=config["data"]["random_state"])

    # 特征缩放
    scaler = StandardScaler()
    x_train_scaled = scaler.fit_transform(x_train)
    x_test_scaled = scaler.transform(x_test)

    return x_train_scaled, x_test_scaled, y_train, y_test

# 预训练数据
def pre_training():
    df = getDef()
    df = addOddEven(df)
    df = getCurrentAvgOddEven(df)
    df = setTarget(df)
    x_train_scaled, x_test_scaled, y_train, y_test = addTrainTest(df)
    return x_train_scaled, x_test_scaled, y_train, y_test
